package org.teremail.mailbox;

import static com.google.common.collect.Maps.newHashMap;
import static java.util.Arrays.asList;
import static java.util.Collections.unmodifiableMap;
import static org.teremail.schema.DBTypes.INT4;
import static org.teremail.schema.DBTypes.TIMESTAMP;
import static org.teremail.schema.DBTypes.VARCHAR;

import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.sql.Timestamp;
import java.sql.Types;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Date;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.UUID;

import javax.sql.DataSource;

import org.teremail.common.Log;
import org.teremail.message.Content;
import org.teremail.message.ContentHeaders;
import org.teremail.message.Entity;
import org.teremail.message.EntityType;
import org.teremail.message.EntityVisitor;
import org.teremail.message.Message;
import org.teremail.message.MessageHeaders;
import org.teremail.message.MultipartEntity;
import org.teremail.message.MultipartMessage;
import org.teremail.message.SimpleEntity;
import org.teremail.message.SimpleMessage;
import org.teremail.message.builder.Builders;
import org.teremail.message.builder.EntityBuilder;
import org.teremail.message.builder.MultipartBuilder;
import org.teremail.message.builder.MultipartEntityBuilder;
import org.teremail.message.builder.MultipartMessageBuilder;
import org.teremail.message.builder.SimpleEntityBuilder;
import org.teremail.message.builder.SimpleMessageBuilder;
import org.teremail.schema.Dialect;
import org.teremail.schema.Schema;
import org.teremail.schema.Table;
import org.teremail.schema.TableBuilder;
import org.teremail.service.Serviceable;
import org.teremail.service.ServiceAdapter;
import org.teremail.service.ServiceRegistrar;
import org.teremail.util.DB;
import org.teremail.util.Formats;

import com.google.common.base.Joiner;
import com.google.common.collect.ArrayListMultimap;
import com.google.common.collect.Multimap;
import com.google.inject.Inject;

/**
 * Database based implementation of the MessageDAO.  It will store a message
 * and all of its parts in a single table.  The table will use the EntityType
 * enum to discriminate between entities of different types.
 * 
 * @author Michael Barker
 *
 */
public class DBMessageDAO implements MessageDAO, Serviceable {

    private final static Log log = Log.getLog(DBMessageDAO.class);
    private final static String UNSPECIFIED = null;
    public final static Table ENTITY_TABLE;

    static {
        TableBuilder entityB = new TableBuilder("ENTITY");
        entityB.column("ID", VARCHAR(40)).primaryKey();
        entityB.column("TYPE", VARCHAR(20));
        entityB.column("CONTENT_HEADERS", VARCHAR(8192)).nullable(true);
        entityB.column("MESSAGE_HEADERS", VARCHAR(10000)).nullable(true);
        entityB.column("CONTENT_ID", VARCHAR(40)).nullable(true);
        entityB.column("PARENT_ID", VARCHAR(40));
        entityB.column("ROOT_ID", VARCHAR(40));
        entityB.column("PART_IDX", INT4());
        entityB.column("SUBJECT", VARCHAR(255)).nullable(true);
        entityB.column("FROM_ADDRESS", VARCHAR(255)).nullable(true);
        entityB.column("SENT_DATE", TIMESTAMP()).nullable(true);
        ENTITY_TABLE = entityB.build();
    }
    
    private final DataSource ds;
    private final Dialect dialect;
    
    @Inject
    public DBMessageDAO(DataSource dataSource, Dialect dialect) {
        this.ds = dataSource;
        this.dialect = dialect;
    }

    public Entity save(Entity entity) {
        SaveVisitor sv = new SaveVisitor();
        entity.accept(sv);
        return sv.getEntityBuilder().build();
    }
    
    private class SaveVisitor implements EntityVisitor {
        
        private final String rootId;
        private final String parentId;
        private EntityBuilder entityBuilder = null;
        private final int partIdx;

        public SaveVisitor(String parentId, String rootId, int partIdx) {
            this.parentId = parentId;
            this.rootId = rootId;
            this.partIdx = partIdx;
        }

        public SaveVisitor() {
            this(UNSPECIFIED, UNSPECIFIED, 0);
        }

        public EntityBuilder getEntityBuilder() {
            return entityBuilder;
        }

        public void visitMultipartEntity(MultipartEntity entity) {
            try {
                String id = UUID.randomUUID().toString();
                String parentId = this.parentId != UNSPECIFIED ? this.parentId : id;
                String rootId = this.rootId != UNSPECIFIED ? this.rootId : id;
                
                insert(id, entity.getType().toString(), 
                        entity.getContentHeader(), 
                        null, null, 
                        parentId, rootId, partIdx);
                
                MultipartEntityBuilder meb = Builders.create(entity);
                meb.setId(id);
                meb.clearParts();
                
                int partIdx = 0;
                for (Entity e : entity) {
                    SaveVisitor sv = new SaveVisitor(parentId, rootId, partIdx++);
                    e.accept(sv);
                    meb.addEntityBuilder(sv.getEntityBuilder());
                }
                
                entityBuilder = meb;
            } catch (SQLException e) {
                throw new RuntimeException(e);
            }
        }

        public void visitMultipartMessage(MultipartMessage message) {
            try {
                String id = UUID.randomUUID().toString();
                String parentId = this.parentId != UNSPECIFIED ? this.parentId : id;
                String rootId = this.rootId != UNSPECIFIED ? this.rootId : id;
                
                insert(id, message.getType().toString(), 
                        message.getContentHeader(), 
                        message.getHeader(), null, 
                        parentId, rootId, partIdx);
                
                MultipartMessageBuilder mmb = Builders.create(message);
                mmb.setId(id);
                mmb.clearParts();
                
                int partIdx = 0;
                for (Entity e : message) {
                    SaveVisitor sv = new SaveVisitor(parentId, rootId, partIdx++);
                    e.accept(sv);
                    mmb.addEntityBuilder(sv.getEntityBuilder());
                }
                
                entityBuilder = mmb;
            } catch (SQLException e) {
                throw new RuntimeException(e);
            }
        }

        public void visitSimpleEntity(SimpleEntity entity) {
            try {
                String id = UUID.randomUUID().toString();
                String parentId = this.parentId != UNSPECIFIED ? this.parentId : id;
                String rootId = this.rootId != UNSPECIFIED ? this.rootId : id;
                
                insert(id, entity.getType().toString(), 
                        entity.getContentHeader(), null, 
                        entity.getContent().getId(), 
                        parentId, rootId, partIdx);
                
                SimpleEntityBuilder seb = Builders.create(entity);
                seb.setId(id);
                entityBuilder = seb;
            } catch (SQLException e) {
                throw new RuntimeException(e);
            }
        }

        public void visitSimpleMessage(SimpleMessage message) {
            try {
                String id = UUID.randomUUID().toString();
                String parentId = this.parentId != UNSPECIFIED ? this.parentId : id;
                String rootId = this.rootId != UNSPECIFIED ? this.rootId : id;
                
                insert(id, message.getType().toString(), 
                        message.getContentHeader(), 
                        message.getHeader(), message.getContent().getId(), 
                        parentId, rootId, partIdx);
                
                SimpleMessageBuilder smb = Builders.create(message);
                smb.setId(id);
                entityBuilder = smb;
            } catch (SQLException e) {
                throw new RuntimeException(e);
            }
        }
    }
    
    /**
     * Inserts an entity of any type
     * 
     * @param id Primary Key.
     * @param type The type of the entity being inserted.
     * @param contentHeaders The ContentHeaders (may be null)
     * @param messageHeaders The MessageHeaders (may be null)
     * @param contentId The contentId, may be -1 to indicate NULL
     * @param parentId The parentId, will be the same as id if this object is
     * the root entity.
     * @param rootId The root message id.
     * 
     * @throws SQLException
     */
    private void insert(String id, String type, ContentHeaders contentHeaders, 
            MessageHeaders messageHeaders, String contentId, String parentId, 
            String rootId, int partIdx) throws SQLException {
        Connection cn = null;
        PreparedStatement ps = null;
        try {
            log.debug("SQL: %s", getInsertStatement());
            cn = ds.getConnection();
            ps = cn.prepareStatement(getInsertStatement());
            ps.setString(1, id);
            ps.setString(2, type);
            setString(ps, 3, contentHeaders);
            setString(ps, 4, messageHeaders);
            setString(ps, 5, contentId);
            ps.setString(6, parentId);
            ps.setString(7, rootId);
            ps.setInt(8, partIdx);
            if (messageHeaders != null) {
                ps.setString(9, messageHeaders.getSubject());
                ps.setString(10, messageHeaders.getFrom());
                
                Date d;
                try {
                    d = Formats.parseSMTPDate(messageHeaders.getDate());
                } catch (Exception e) {
                    d = new Date();
                }
                ps.setTimestamp(11, new Timestamp(d.getTime()));
                
            } else {
                ps.setNull(9, Types.VARCHAR);
                ps.setNull(10, Types.VARCHAR);
                ps.setNull(11, Types.TIMESTAMP);
            }
            ps.executeUpdate();
        } finally {
            DB.close(ps, cn);
        }
    }
    
    /**
     * Read an entityBuilder from the tuple that was read from
     * the database.
     * 
     * @param id
     * @param type
     * @param contentHeaders
     * @param messageHeaders
     * @param contentId
     * @return
     */
    private EntityBuilder read(String id, String type, String contentHeaders, 
            String messageHeaders, String contentId) {
                
        EntityType entityType = EntityType.valueOf(type);
        
        switch (entityType) {
        case SIMPLE_MESSAGE:
            SimpleMessageBuilder smb = new SimpleMessageBuilder();
            smb.setId(id);
            smb.setHeaders(messageHeaders);
            smb.setContent(new Content(contentId));
            return smb;
            
        case SIMPLE_ENTITY:
            SimpleEntityBuilder seb = new SimpleEntityBuilder();
            seb.setId(id);
            seb.setContent(new Content(contentId));
            seb.setContentHeaders(contentHeaders);
            return seb;
            
        case MULTIPART_MESSAGE:
            MultipartMessageBuilder mmb = new MultipartMessageBuilder();
            mmb.setId(id);
            mmb.setHeaders(messageHeaders);
            mmb.setContentHeaders(contentHeaders);
            return mmb;
            
        case MULTIPART_ENTITY:
            MultipartEntityBuilder meb = new MultipartEntityBuilder();
            meb.setId(id);
            meb.setContentHeaders(contentHeaders);
            return meb;
            
            
        default:
            throw new RuntimeException("Unknown entity type: " + entityType);
        }
    }
    
    /**
     * @see org.teremail.mailbox.MessageDAO#load(java.util.Collection)
     */
    public Map<String, Message> loadMessages(Collection<String> ids) {
        Map<String, Entity> entities = load(ids);
        Map<String, Message> messages = newHashMap();
        for (Entity e : entities.values()) {
            if (e instanceof Message) {
                messages.put(e.getId(), (Message) e);
            }
        }
        return unmodifiableMap(messages);
    }
    
    
    /**
     * @see org.teremail.mailbox.MessageDAO#load(java.util.Collection)
     */
    public Map<String, Entity> load(Collection<String> ids) {
        Connection cn = null;
        PreparedStatement ps = null;
        ResultSet rs = null;
        
        try {
            cn = ds.getConnection();
            ps = cn.prepareStatement(getSelectInList(ids));
            int i = 1;
            for (String id : ids) {
                ps.setString(i++, id);
            }
            
            rs = ps.executeQuery();
            Multimap<String,EntityBuilder> builders = ArrayListMultimap.create();
            
            List<EntityBuilder> roots = new ArrayList<EntityBuilder>();
            Set<String> idSet = new HashSet<String>(ids);
            
            while (rs.next()) {
                String id = rs.getString(1);
                String type = rs.getString(2);
                String contentHeaders = rs.getString(3);
                String messageHeaders = rs.getString(4);
                String contentId = rs.getString(5);
                String parentId = rs.getString(6);
                
                EntityBuilder eb = read(id, type, contentHeaders, 
                        messageHeaders, contentId);
                
                if (idSet.contains(eb.getId())) {
                    // Store the root object explicitly
                    roots.add(eb);
                } else {
                    // Store all other indexed by the parent id.
                    builders.put(parentId, eb);
                }
                
            } 
            
            Map<String,Entity> messages = newHashMap();
            for (EntityBuilder root : roots) {
                build(root, builders);
                Entity e = root.build();
                messages.put(root.getId(), e);
            }
            
            return unmodifiableMap(messages);
                        
        } catch (SQLException e) {
            throw new RuntimeException(e);
        } finally {
            DB.close(rs, ps, cn);
        }
    }

    public Entity load(String rootId) {
        Map<String,Entity> m = load(asList(rootId));
            
        if (!m.containsKey(rootId)) {
            throw new RuntimeException("Unable to find message with id: " + rootId);
        }
        
        return m.get(rootId);
    }
    
    public Message loadMessage(String rootId) {
        return (Message) load(rootId);
    }

    /**
     * Rebuilds an entity from the flat structure into a heirachical one.
     * 
     * @param root
     * @param builders
     * @return
     */
    private void build(EntityBuilder parent, Multimap<String, EntityBuilder> builders) {
        
        if (parent.isMultipart()) {
            MultipartBuilder mb  = (MultipartBuilder) parent;
            Collection<EntityBuilder> ebs = builders.get(mb.getId());
            for (EntityBuilder eb : ebs) {
                build(eb, builders);
                mb.addEntityBuilder(eb);
            }
        }
    }

    private static void setString(PreparedStatement ps, int idx, Object o) 
            throws SQLException {
        if (o == null) {
            ps.setNull(idx, Types.VARCHAR);
        } else {
            ps.setString(idx, o.toString());
        }
    }
    
    private String getInsertStatement() {
        return "INSERT INTO entity (id, type, content_headers, message_headers, content_id, parent_id, root_id, part_idx, subject, from_address, sent_date) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)";
    }

    private String getSelectInList(Collection<String> ids) {
        
        String[] s = new String[ids.size()];
        for (int i = 0; i < s.length; i++) {
            s[i] = "?";
        }
        String in = "(" + Joiner.on(", ").join(s) + ")";
        
        return "SELECT id," +
                "      type," +
                "      content_headers," +
                "      message_headers," +
                "      content_id," +
                "      parent_id," +
                "      root_id " +
                "FROM  entity " +
                "WHERE root_id IN " + in + " " +
                "ORDER BY part_idx";
    }

    // ==================
    // Service Management
    // ==================
    
    public void init() {
        Schema.createIfNotExists(ds, ENTITY_TABLE, dialect);
    }
    
    @Inject
    public void registerService(ServiceRegistrar serviceRegistrar) {
        serviceRegistrar.register(new ServiceAdapter() {
            public void start() {
                init();
            }
        });
    }
}
